package com.unitd.frame.mybatis.datasource;

import com.unitd.frame.mybatis.plugin.MybatisPluginInterceptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @desc 读取和设置数据源信息的上下文环境的类,主要负责注册、改变上下文数据源
 * @filename DataSourceContextHolder.java
 * @copyright www.unitd.com
 * @author Hudan
 * @version 1.0
 * @date 2016/12/17
 */
public class DataSourceContextHolder {

	private static final Logger logger = LoggerFactory.getLogger(DataSourceContextHolder.class);

	/** 多线程环境下,非线程安全的基本数据类型会引起 ,故使用线程安全的基本数据类型 */
	private final static AtomicLong counter = new AtomicLong(10);
	/** 主库集合 */
	private static final Map<String, String> masters = new HashMap<>();
	/** 从库集合 */
	private static final Map<String, List<String>> slaves = new HashMap<>();
	/** 数据源 */
	private final ThreadLocal<DataSourceContextVals> contextVals = new ThreadLocal<>();
	/** 数据源信息的上下文持有器 */
	private static volatile DataSourceContextHolder holder = new DataSourceContextHolder();

	private DataSourceContextHolder() {
	}

	public static DataSourceContextHolder get() {
		return holder;
	}

	/**
	 * @desc 根据数据源对象bean的key向数据源集合中注册数据源对象,方便后续使用的时候直接根据key获取
	 * @param dsKey 数据源bena的key
	 */
	protected void registerDataSourceKey(String dsKey) {
		String dbIndex = "0";
		if (dsKey.startsWith("group")) {
			dbIndex = dsKey.split("\\_")[0].replace("group", "");
		}
		if (dsKey.contains("master")) {
			masters.put(dbIndex, dsKey);
		} else {
			List<String> sameDbSlaves;
			if (slaves.containsKey(dbIndex)) {
				sameDbSlaves = slaves.get(dbIndex);
			} else {
				sameDbSlaves = new ArrayList<>();
				slaves.put(dbIndex, sameDbSlaves);
			}
			sameDbSlaves.add(dsKey);
		}
	}

	/**
	 * @desc 设置分库使用数据库序列(groupId)
	 * @param dbIndex 分库数据库序列
	 */
	public void setDbIndex(int dbIndex) {
		DataSourceContextVals vals = contextVals.get();
		if (vals == null) {
			vals = new DataSourceContextVals();
		}

		vals.dbIndex = dbIndex;
		vals.routedBeforeGetConn = true;
		contextVals.set(vals);
	}

	/**
	 * @desc 设置是否使用从库
	 * @param useSlave 是否使用从库
	 * @return 从库上下文持有器
	 */
	public DataSourceContextHolder useSlave(boolean useSlave) {
		DataSourceContextVals vals = contextVals.get();
		if (vals == null) {
			vals = new DataSourceContextVals();
		}

		vals.userSlave = useSlave;
		vals.routedBeforeGetConn = true;
		contextVals.set(vals);
		return this;
	}

	/**
	 * @desc 获取当前数据源名
	 * @return 当前数据源名
	 */
	protected String getDataSourceKey() {
		if(!MybatisPluginInterceptor.isRwRouteEnabled() && !MybatisPluginInterceptor.isDbShardEnabled()){
			return masters.get(0);
		}

		DataSourceContextVals vals = contextVals.get();
		if (vals == null) {
			vals = new DataSourceContextVals();
			vals.forceMaster = true;
			contextVals.set(vals);

			return masters.get(0);
		}

		int dbGoupId = vals.dbIndex;
		String dsKey;

		if (vals.forceMaster || !vals.userSlave || !vals.routedBeforeGetConn) {
			if (dbGoupId > 0 && masters.size() <= dbGoupId + 1) {
				throw new RuntimeException("expect db group number is :" + dbGoupId + ",actaul:" + (dbGoupId + 1));
			}
			dsKey = masters.get(String.valueOf(dbGoupId));
		} else {
			if (dbGoupId > 0 && slaves.size() <= dbGoupId + 1) {
				throw new RuntimeException("expect db group number is :" + dbGoupId + ",actaul:" + (dbGoupId + 1));
			}
			dsKey = selectSlave(dbGoupId);
		}

		vals.dsKey = dsKey;
		logger.debug("current route rule is:userSlave[{}]|forceMaster[{}]|routedBeforeGetConn[{}], use dataSource key is [{}]!", vals.userSlave, vals.forceMaster, vals.routedBeforeGetConn, vals.dsKey);

		// 重置路由状态
		vals.routedBeforeGetConn = false;

		return dsKey;
	}

	/**
	 * @desc 判断是否强制使用一种方式
	 * @return boolean
	 */
	public boolean isForceUseMaster() {
		DataSourceContextVals vals = contextVals.get();
		if (vals == null)
			return false;
		return vals.forceMaster;
	}

	/**
	 * @desc 轮循分配slave节点
	 * @return String
	 */
	private static String selectSlave(Serializable dbIndex) {
		List<String> sameDbSlaves = slaves.get(dbIndex.toString());
		// 无从库则路由到主库
		if (sameDbSlaves == null || sameDbSlaves.isEmpty()) {
			String masterKey = masters.get(dbIndex.toString());
			logger.debug("current no slave found ,default use [{}]!", masterKey);
			return masterKey;
		}
		if (sameDbSlaves.size() == 1)
			return sameDbSlaves.get(0);

		// 获取从库的索引,并返回
		int selectIndex = (int) (counter.getAndIncrement() % sameDbSlaves.size());
		return sameDbSlaves.get(selectIndex);
	}

	public void clear() {
		contextVals.remove();
	}

	/**
	 * @desc 数据源上下文对象值集合对象类
	 */
	private class DataSourceContextVals {
		/** 分库的数据库序列 */
		private int dbIndex;
		/** 是否使用从库 */
		private boolean userSlave;
		/** 是否强制使用主库 */
		private boolean forceMaster;
		/** 数据源的key */
		private String dsKey;
		/** 标记在获取连接前是否执行了路由,在事务的情况获取连接在路由之前 */
		private boolean routedBeforeGetConn;
	}
}